package visited

import "sync"

type VisitedType struct {
	mu    sync.RWMutex
	store map[string][2]int
}

func NewVisitedType() *VisitedType {
	return &VisitedType{
		store: map[string][2]int{},
	}
}

func (v *VisitedType) Get(key string) (result bool) {
	v.mu.RLock()
	defer v.mu.RUnlock()

	if value, ok := v.store[key]; ok {
		if value[0] == 1 {
			return true
		} else {
			return false
		}
	}
	return
}

func (v *VisitedType) Set(key string, value bool) {
	v.mu.Lock()
	defer v.mu.Unlock()

	i := 0
	if value == true {
		i = 1
	}
	if item, ok := v.store[key]; ok {
		item[0] = i
		v.store[key] = item
	} else {
		v.store[key] = [2]int{i, 0}
	}
}

func (v *VisitedType) SetFlagAndAddCount(key string, value bool, i int) {
	v.mu.Lock()
	defer v.mu.Unlock()

	flag := 0
	if value == true {
		flag = 1
	}
	if item, ok := v.store[key]; ok {
		item[0] = flag
		item[1] += i
		v.store[key] = item
	} else {
		v.store[key] = [2]int{flag, i}
	}
}

func (v *VisitedType) GetCount(key string) (result int) {
	v.mu.RLock()
	defer v.mu.RUnlock()

	if value, ok := v.store[key]; ok {
		return value[1]
	}
	return
}

func (v *VisitedType) AddCount(key string, i int) {
	v.mu.Lock()
	defer v.mu.Unlock()

	if item, ok := v.store[key]; ok {
		item[1] += i
		v.store[key] = item
	}
}

func (v *VisitedType) Clear() {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.store = map[string][2]int{}
}
